#include "vm/swap.h"
#include "devices/block.h"
#include "threads/vaddr.h"
#include "threads/synch.h"
#include "bitmap.h"

#define SECTORS_PER_PAGE (PGSIZE / BLOCK_SECTOR_SIZE)

static struct block *swap_block;    // The swap block device.
static struct bitmap *swap_table;  // Bitmap to track free swap slots.
static struct lock swap_lock;      // Lock to synchronize access to the swap table.

/* Initializes the swap table. */
void swap_table_init(void) {
    swap_block = block_get_role(BLOCK_SWAP);
    if (swap_block == NULL)
        PANIC("No swap partition found!");

    size_t swap_size = block_size(swap_block) / SECTORS_PER_PAGE;
    swap_table = bitmap_create(swap_size);
    if (swap_table == NULL)
        PANIC("Failed to create swap table!");

    lock_init(&swap_lock);
}

/* Allocates a swap slot. */
size_t swap_allocate(void) {
    lock_acquire(&swap_lock);
    size_t idx = bitmap_scan_and_flip(swap_table, 0, 1, false);
    if (idx == BITMAP_ERROR)
        PANIC("Swap partition is full!");
    lock_release(&swap_lock);
    return idx;
}

/* Frees a swap slot. */
void swap_free(size_t idx) {
    lock_acquire(&swap_lock);
    bitmap_reset(swap_table, idx);
    lock_release(&swap_lock);
}

/* Swaps out a page to the swap partition. */
void swap_out(void *kpage) {
    size_t idx = swap_allocate();
    for (size_t i = 0; i < SECTORS_PER_PAGE; i++) {
        block_write(swap_block, idx * SECTORS_PER_PAGE + i,
                    (uint8_t *)kpage + i * BLOCK_SECTOR_SIZE);
    }
}

/* Swaps in a page from the swap partition. */
void swap_in(size_t idx, void *kpage) {
    for (size_t i = 0; i < SECTORS_PER_PAGE; i++) {
        block_read(swap_block, idx * SECTORS_PER_PAGE + i,
                   (uint8_t *)kpage + i * BLOCK_SECTOR_SIZE);
    }
    swap_free(idx);
}
